#include <stdint.h>
#include <string.h>
#include "stm32f4xx.h"

#define STM32_FLASH_SIZE (1024 * 1024) // 1MB

// Flash sector start addresses for STM32F407
static const uint32_t sector_addr_begin[] = {
    0x08000000, // Sector 0, 16 KB
    0x08004000, // Sector 1, 16 KB
    0x08008000, // Sector 2, 16 KB
    0x0800C000, // Sector 3, 16 KB
    0x08010000, // Sector 4, 64 KB
    0x08020000, // Sector 5, 128 KB
    0x08040000, // Sector 6, 128 KB
    0x08060000, // Sector 7, 128 KB
    0x08080000, // Sector 8, 128 KB
    0x080A0000, // Sector 9, 128 KB
    0x080C0000, // Sector 10, 128 KB
    0x080E0000, // Sector 11, 128 KB
};

/**
 * @brief Get flash sector number for a given address
 * @param addr: Address in flash memory
 * @return Sector number or -1 if invalid address
 */
int ota_flash_get_sector(uint32_t addr)
{
    int sector = -1;

    // Check if address is within valid flash range
    if (addr < sector_addr_begin[0] || addr >= (sector_addr_begin[0] + STM32_FLASH_SIZE)) {
        return sector; // Invalid address
    }

    // Find the sector containing the address
    for (uint32_t i = 0; i < sizeof(sector_addr_begin) / sizeof(sector_addr_begin[0]); i++) {
        if (addr >= sector_addr_begin[i]) {
            sector = i;
        } else {
            break;
        }
    }
    return sector;
}

/**
 * @brief Check if address is aligned to sector boundary
 * @param addr: Address to check
 * @return 1 if aligned, 0 if not aligned
 */
int ota_flash_check_align(uint32_t addr)
{
    for (uint32_t i = 0; i < sizeof(sector_addr_begin) / sizeof(sector_addr_begin[0]); i++) {
        if (addr == sector_addr_begin[i]) {
            return 1;
        }
    }
    return 0;
}

/**
 * @brief Erase flash sectors
 * @param addr: Start address
 * @param length: Number of bytes to erase
 * @return 0 on success, -1 on error
 */
int ota_flash_erase(uint32_t addr, uint32_t length)
{
    int ret = 0;
    int first_sector = ota_flash_get_sector(addr);
    int last_sector = ota_flash_get_sector(addr + length - 1);

    // Validate sector range
    if (first_sector == -1 || last_sector == -1 || last_sector < first_sector) {
        return -1; // Invalid sector
    }

    FLASH_Unlock();
    // Clear all pending flags
    FLASH_ClearFlag(FLASH_FLAG_EOP | FLASH_FLAG_OPERR | FLASH_FLAG_WRPERR |
                   FLASH_FLAG_PGAERR | FLASH_FLAG_PGPERR | FLASH_FLAG_PGSERR);

    // Erase each sector in the range
    for (int sector = first_sector; sector <= last_sector; sector++) {
        FLASH_Status status = FLASH_EraseSector(sector * 8, VoltageRange_3); // Sector * 8 for SPL sector numbering
        if (status != FLASH_COMPLETE) {
            ret = -2;
            break;
        }
    }
    FLASH_Lock();

    return ret;
}

/**
 * @brief Write data to flash memory
 * @param addr: Destination address in flash
 * @param data: Source data pointer
 * @param len: Number of bytes to write
 * @return 0 on success, -1 on error
 */
int ota_flash_write(uint32_t addr, const uint8_t* data, uint32_t len)
{
    uint32_t index, head_fill_len, last_fill_len, program_word_len, program_word;
    const uint8_t *dst_data, *src_data;

    dst_data = (const uint8_t *)addr;
    src_data = data;

    // Calculate head alignment fill length
    head_fill_len = addr & 0x3;
    head_fill_len = 4 - head_fill_len;
    head_fill_len &= 0x3;

    // Calculate program lengths based on alignment and total length
    if (len <= head_fill_len) {
        head_fill_len = len;
        program_word_len = 0;
        last_fill_len = 0;
    } else {
        program_word_len = (len - head_fill_len) / 4;
        last_fill_len = (len - head_fill_len) & 0x3;
    }

    FLASH_Unlock();

    // Clear all pending flags
    FLASH_ClearFlag(FLASH_FLAG_EOP | FLASH_FLAG_OPERR | FLASH_FLAG_WRPERR |
                   FLASH_FLAG_PGAERR | FLASH_FLAG_PGPERR | FLASH_FLAG_PGSERR);

    // Write head unaligned bytes (byte programming)
    for (index = 0; index < head_fill_len; index++) {
        if (FLASH_ProgramByte(addr, *data) != FLASH_COMPLETE) {
            FLASH_Lock();
            return -1;
        }
        addr++;
        data++;
    }

    // Write aligned word data (word programming)
    for (index = 0; index < program_word_len; index++) {
        memcpy(&program_word, data, 4);
        if (FLASH_ProgramWord(addr, program_word) != FLASH_COMPLETE) {
            FLASH_Lock();
            return -2;
        }
        addr += 4;
        data += 4;
    }

    // Write remaining unaligned bytes (byte programming)
    for (index = 0; index < last_fill_len; index++) {
        if (FLASH_ProgramByte(addr, *data) != FLASH_COMPLETE) {
            FLASH_Lock();
            return -3;
        }
        addr++;
        data++;
    }

    FLASH_Lock();

    // Verify written data
    for (index = 0; index < len; index++) {
        if (dst_data[index] != src_data[index]) {
            return -4; // Verification failed
        }
    }

    return 0;
}
